"""
Example: 2D Helmholtz Equation

Reproduces results from Section 4.2 of the paper.
Solves: ∂²u/∂x² + ∂²u/∂y² - λu = f(x,y) on [0,3.6]×[0,3.6] with Dirichlet BCs.

Exact solution: u = -[3/2*cos(πx+2π/5) + 2*cos(2πx-π/5)] * [3/2*cos(πy+2π/5) + 2*cos(2πy-π/5)]
"""

import numpy as np
import sys
import os

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from locELM.core.domain import Domain
from locELM.core.networks import MultiSubdomainNetwork
from locELM.solvers.linear import LinearPDESolver, helmholtz_operator_2d
from locELM.utils.helpers import (
    compute_errors, plot_solution_2d, Timer, print_results,
    create_evaluation_function
)


def exact_solution(points):
    """Exact solution to the 2D Helmholtz equation."""
    if points.ndim == 1:
        points = points.reshape(1, -1)

    x = points[:, 0]
    y = points[:, 1]

    term_x = 1.5 * np.cos(np.pi*x + 2*np.pi/5) + 2 * np.cos(2*np.pi*x - np.pi/5)
    term_y = 1.5 * np.cos(np.pi*y + 2*np.pi/5) + 2 * np.cos(2*np.pi*y - np.pi/5)

    u = -term_x * term_y
    return u.reshape(-1, 1) if u.ndim == 1 else u.reshape(-1, 1)


def source_term(points):
    """Source term f(x,y) corresponding to exact solution."""
    lam = 10.0
    u = exact_solution(points)

    if points.ndim == 1:
        points = points.reshape(1, -1)

    x = points[:, 0]
    y = points[:, 1]

    # Analytical second derivatives
    term_x = 1.5 * np.cos(np.pi*x + 2*np.pi/5) + 2 * np.cos(2*np.pi*x - np.pi/5)
    term_y = 1.5 * np.cos(np.pi*y + 2*np.pi/5) + 2 * np.cos(2*np.pi*y - np.pi/5)

    d2_term_x = (-1.5*np.pi**2 * np.cos(np.pi*x + 2*np.pi/5)
                 - 8*np.pi**2 * np.cos(2*np.pi*x - np.pi/5))
    d2_term_y = (-1.5*np.pi**2 * np.cos(np.pi*y + 2*np.pi/5)
                 - 8*np.pi**2 * np.cos(2*np.pi*y - np.pi/5))

    d2u_dx2 = -d2_term_x * term_y
    d2u_dy2 = -term_x * d2_term_y

    laplacian_u = d2u_dx2 + d2u_dy2
    f = laplacian_u - lam * u.flatten()

    return f.flatten()


def run_helmholtz_2d(n_subdomains_x=2, n_subdomains_y=2,
                     n_colloc_x=25, n_colloc_y=25,
                     n_params_per_subdomain=400, r_max=1.5, seed=1):
    """
    Run 2D Helmholtz equation example.

    Parameters
    ----------
    n_subdomains_x : int
        Number of subdomains in x direction
    n_subdomains_y : int
        Number of subdomains in y direction
    n_colloc_x : int
        Number of collocation points in x per subdomain
    n_colloc_y : int
        Number of collocation points in y per subdomain
    n_params_per_subdomain : int
        Number of training parameters (M) per subdomain
    r_max : float
        Maximum magnitude of random coefficients
    seed : int
        Random seed

    Returns
    -------
    dict
        Results dictionary
    """
    n_subdomains = n_subdomains_x * n_subdomains_y

    print(f"\n{'='*70}")
    print(f"2D Helmholtz Equation - locELM Method")
    print(f"{'='*70}")
    print(f"Domain: [0, 3.6] × [0, 3.6]")
    print(f"Subdomains: {n_subdomains_x} × {n_subdomains_y} = {n_subdomains}")
    print(f"Collocation points per subdomain: {n_colloc_x} × {n_colloc_y}")
    print(f"Training parameters per subdomain: {n_params_per_subdomain}")
    print(f"Random coefficient magnitude (Rm): {r_max}")
    print(f"{'='*70}\n")

    # Problem parameters
    lam = 10.0
    a1, b1 = 0.0, 3.6
    a2, b2 = 0.0, 3.6

    # Define domain
    domain = Domain(
        bounds=[(a1, b1), (a2, b2)],
        n_subdomains=[n_subdomains_x, n_subdomains_y]
    )

    # Create neural networks
    networks = MultiSubdomainNetwork(
        n_subdomains=n_subdomains,
        input_dim=2,
        output_dim=1,
        hidden_layers=[n_params_per_subdomain],  # One hidden layer
        activation='tanh',
        r_max=r_max,
        seed=seed
    )

    # Add networks for each subdomain
    for idx in range(n_subdomains):
        region = domain.get_subdomain_region(idx)
        networks.add_subdomain_network(region)

    # Create solver
    solver = LinearPDESolver(
        domain=domain,
        networks=networks,
        continuity_order=[1, 1],  # C^1 continuity in both x and y
        collocation_type='uniform',
        n_collocation_points=[n_colloc_x, n_colloc_y]
    )

    # Define PDE operator
    def pde_operator(network, points, subdomain_idx):
        return helmholtz_operator_2d(network, points, subdomain_idx, lam=lam)

    # Define boundary conditions
    def bc_lower_x(points):
        x_val = np.full((points.shape[0], 2), a1)
        x_val[:, 1] = points[:, 1]
        return exact_solution(x_val).flatten()

    def bc_upper_x(points):
        x_val = np.full((points.shape[0], 2), b1)
        x_val[:, 1] = points[:, 1]
        return exact_solution(x_val).flatten()

    def bc_lower_y(points):
        y_val = points.copy()
        y_val[:, 1] = a2
        return exact_solution(y_val).flatten()

    def bc_upper_y(points):
        y_val = points.copy()
        y_val[:, 1] = b2
        return exact_solution(y_val).flatten()

    boundary_conditions = {
        0: {  # x-dimension
            'lower': bc_lower_x,
            'upper': bc_upper_x
        },
        1: {  # y-dimension
            'lower': bc_lower_y,
            'upper': bc_upper_y
        }
    }

    # Solve
    print("Solving PDE...")
    with Timer() as timer:
        weights = solver.solve_steady_state(
            pde_operator=pde_operator,
            source_term=source_term,
            boundary_conditions=boundary_conditions,
            output_dim=1
        )

    training_time = timer.elapsed_time
    print(f"Training completed in {training_time:.4f} seconds")

    # Evaluate solution
    eval_func = create_evaluation_function(networks, domain, output_dim=1)

    # Compute errors
    errors = compute_errors(
        solution_func=eval_func,
        exact_func=lambda p: exact_solution(p).flatten(),
        domain_bounds=[(a1, b1), (a2, b2)],
        n_eval_points=[100, 100]
    )

    # Print results
    print_results(errors, training_time, method_name="locELM (2D Helmholtz)")

    # Plot solution
    plot_solution_2d(
        solution_func=eval_func,
        domain_bounds=[(a1, b1), (a2, b2)],
        exact_func=lambda p: exact_solution(p).flatten(),
        n_points=[100, 100],
        title='2D Helmholtz Equation Solution',
        save_path='helmholtz_2d_solution.png'
    )
    print("Solution plot saved to 'helmholtz_2d_solution.png'")

    return {
        'errors': errors,
        'timing': training_time,
        'weights': weights,
        'networks': networks,
        'domain': domain
    }


if __name__ == "__main__":
    # Reproduce paper results
    print("\n" + "="*70)
    print("Reproducing Paper Results - 2D Helmholtz Equation")
    print("="*70)

    # Test 1: Effect of number of subdomains
    print("\n### Test 1: Effect of number of subdomains ###")
    for nx in [1, 2, 3]:
        results = run_helmholtz_2d(
            n_subdomains_x=nx,
            n_subdomains_y=nx,
            n_colloc_x=25,
            n_colloc_y=25,
            n_params_per_subdomain=400,
            r_max=1.5,
            seed=1
        )
        print(f"Grid: {nx}×{nx}, Max Error = {results['errors']['max_error']:.6e}, "
              f"Time = {results['timing']:.4f}s")

    # Test 2: Effect of collocation points (4 subdomains)
    print("\n### Test 2: Effect of collocation points ###")
    for n_colloc in [15, 20, 25, 30]:
        results = run_helmholtz_2d(
            n_subdomains_x=2,
            n_subdomains_y=2,
            n_colloc_x=n_colloc,
            n_colloc_y=n_colloc,
            n_params_per_subdomain=400,
            r_max=1.5,
            seed=1
        )
        print(f"Colloc: {n_colloc}×{n_colloc}, Max Error = {results['errors']['max_error']:.6e}, "
              f"Time = {results['timing']:.4f}s")

    # Test 3: High accuracy configuration
    print("\n### Test 3: High accuracy configuration ###")
    results = run_helmholtz_2d(
        n_subdomains_x=2,
        n_subdomains_y=2,
        n_colloc_x=30,
        n_colloc_y=30,
        n_params_per_subdomain=600,
        r_max=1.5,
        seed=1
    )
